import os
import sys
import json
import numpy as np


def get_bitrate_count(movie):
    """
    获取视频比特率等级数量
    :param movie:
    :return:
    """
    bitrate_level = 0
    if movie == 'CMPVP907':
        bitrate_level = 6
    elif movie == 'skiing':
        bitrate_level = 4
    else:
        for file in os.listdir(basePath + movie + '/face0/video/vp09'):
            file_path = os.path.join(basePath + movie + '/face0/video/vp09', file)
            if os.path.isdir(file_path):
                bitrate_level += 1
    return bitrate_level


def get_segment_count(movie):
    """
    获取视频段的数量
    :param movie:
    :return:
    """
    if movie == 'CMPVP907':
        files = os.listdir(basePath + movie + '/face0/output/video/vp09/1')
    else:
        files = os.listdir(basePath + movie + '/face0/video/vp09/1')
    return len(files) - 1


def get_size_noTiled(movie):
    """
    针对 6*1*1分片下的视频信息
    :param movie:
    :return:
    """
    path = '../' + movie + '/face'
    path = basePath + movie + '/face'
    bitrate_count = get_bitrate_count(movie)
    segment_count = get_segment_count(movie)
    file_size = np.zeros((segment_count, 6, bitrate_count))

    for i in range(6):
        file_path = path + str(i)
        if movie == 'CMPVP907':
            file_path = file_path + '/output'
        for j in range(bitrate_count):
            # file_list = os.listdir(os.path.join(file_path, 'video/vp09/' + str(j)))
            for k in range(segment_count):
                file_size[k][i][j] = os.path.getsize(
                    file_path + '/video/vp09/' + str(j + 1) + '/seg-' + str(k + 1) + '.m4s')  # Byte

    bitrates = np.round(np.mean(np.mean(file_size, axis=0), axis=0) * 8, decimals=2)  # bps
    per_tile_bitrates = np.round(np.mean(file_size, axis=0) * 8, decimals=2)  # bps
    # print(per_tile_bitrates.T.tolist())

    dic = {'segment_duration_ms': 1000, 'bitrate_bps': bitrates.tolist(),
           'bitrate_per_tile_bps': per_tile_bitrates.T.tolist(), 'segment_size_bytes': file_size.tolist()}
    # tile级别的平均比特率等级
    # 不同比特率等级中每个tile平均比特率等级
    # 每个segment中，每个tile的不同质量
    with open(basePath + movie + '/' + movie + '.json', 'w') as f:
        json.dump(dic, f, indent=4)


def get_counts(movie):
    """
    返回分片视频统计信息
    :param movie:
    :return:每个面的tile数量、比特率等级、视频段数量
    """
    return len(os.listdir(basePath + movie + '/face0/')), \
           len(os.listdir(basePath + movie + '/face0/tile1/video/vp09/')), \
           len(os.listdir(basePath + movie + '/face0/tile1/video/vp09/1/')) - 1


def get_size_tiled(movie):
    """
    针对 6*n*n分片下的视频信息
    :param movie:
    :return:
    """
    # path = '../' + movie + '/face'
    path = basePath + movie + '/face'
    tile_count, bitrate_count, segment_count = get_counts(movie)
    file_size = np.zeros((segment_count, 6, tile_count, bitrate_count))
    for i in range(6):
        file_path = path + str(i) + ''
        for j in range(tile_count):
            for m in range(bitrate_count):
                for n in range(segment_count):
                    file_size[n][i][j][m] = os.path.getsize(
                        file_path + '/tile' + str(j + 1) + '/video/vp09/' + str(m + 1)
                        + '/seg-' + str(n + 1) + '.m4s')

    bitrates = np.round(np.mean(np.mean(np.mean(file_size, axis=0), axis=0), axis=0), decimals=2)
    per_tile_bitrates = np.round(np.mean(file_size, axis=0), decimals=2).swapaxes(1, 2).swapaxes(0, 1)

    dic = {'segment_duration_ms': 1000, 'bitrate_bps': bitrates.tolist(),
           'bitrate_per_tile_bps': per_tile_bitrates.tolist(), 'segment_size_bytes': file_size.tolist()}
    # tile级别的平均比特率等级
    # bitrate levels*4 -> Face*6 -> tile per face*4
    # Segment -> Face*6 -> tile per face*4 -> bitrate levels
    with open(basePath + movie + '/' + movie + '.json', 'w') as f:
        json.dump(dic, f, indent=4)


if __name__ == '__main__':
    basePath = '../../'
    movie = sys.argv[1]
    if not os.path.isdir(basePath + movie):
        print('Movie not exist!')

    else:
        if movie == 'tiledSkiing' or movie == 'tiledSkiing_old':
            get_size_tiled(movie)
            print('Successfully processed ' + movie)
        # 6 * 1 * 1
        elif os.path.isdir(basePath + movie + '/face0'):
            get_size_noTiled(movie)
            print('Successfully processed ' + movie)
        # 6 * n * n
        elif os.path.isdir(basePath + movie + '/face0_tiled'):
            get_size_tiled(movie)
            print('Successfully processed ' + movie)
        else:
            print('Nothing happened')
